package org.shx.springboot.ollama.ai.service.impl;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.node.ArrayNode;
import com.fasterxml.jackson.databind.node.ObjectNode;
import jakarta.annotation.Resource;
import jakarta.servlet.http.HttpSession;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.shx.springboot.ollama.ai.controller.ChatController;
import org.shx.springboot.ollama.ai.service.ChatService;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.chat.prompt.PromptTemplate;
import org.springframework.ai.document.Document;
import org.springframework.ai.ollama.OllamaChatModel;
import org.springframework.ai.tool.ToolCallback;
import org.springframework.ai.tool.ToolCallbackProvider;
import org.springframework.ai.tool.definition.ToolDefinition;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.model.Generation;
import org.springframework.http.MediaType;
import org.springframework.stereotype.Service;
import org.springframework.web.reactive.function.client.WebClient;
import reactor.core.publisher.Flux;

import java.util.ArrayList;
import java.util.List;
import java.util.stream.Collectors;

/*知识库相关*/
@Slf4j
@Service
@RequiredArgsConstructor
public class ChatServiceImpl implements ChatService {

    @Resource
    @Qualifier("ollamaChatClient")
    private ChatClient ollamaChatClient;
    @Autowired
    private OllamaChatModel chatModel;
    @Resource
    private VectorStore vectorStore;
    @Resource
    private ToolCallbackProvider callbackProvider;


    @Override
    public String chat(String message) {
        return chatModel.call(message); // OllamaApi类
    }

    @Override
    public Flux<ChatResponse> streamChat(String message) {
        Prompt prompt = new Prompt(new UserMessage(message));
        log.info("开始调用OllamaChatModel的stream方法"); // 添加此日志
//        List<Message> messageList = List.of(new UserMessage(message), new SystemMessage("中文回答"));
//        Prompt prompt = new Prompt(messageList);
        return chatModel.stream(prompt);
    }

    @Override
    public Flux<String> generateStreamChat(String message) {
        List<Message> messageList = List.of(
                new UserMessage(message),
                new SystemMessage("中文回答")
        );
        Prompt prompt = new Prompt(messageList);

        Flux<String> sseFlux = chatModel.stream(prompt)
                .map(chatResponse -> {
                    // 确保正确提取文本（根据实际结构调整，避免空值）
                    String content = chatResponse.getResult().getOutput().getText();
//                    String sseContent = "data: " + (content != null ? content : "") + "\n\n";
                    return content;
                });

        // 显式设置 Content-Type，强化匹配
        return sseFlux;
    }

    // 最大历史消息数量（防止内存溢出）
    private static final int MAX_HISTORY_SIZE = 20;
    @Override
    public Flux<String> generateStreamChatSession(String message, HttpSession session) {

        log.info("开始调用OllamaChatModel的generateStreamChatSession方法"); // 添加此日志

        // 1. 从会话中获取历史消息，若不存在则初始化
        List<Message> messageList = (List<Message>) session.getAttribute(ChatController.CHAT_HISTORY_KEY);
        if (messageList == null) {
            messageList = new ArrayList<>();
            // 添加系统提示（仅首次初始化时添加）
            messageList.add(new SystemMessage("中文回答，保持对话连贯，记住上下文信息"));
        }

        // 2. 添加当前用户消息到历史
        messageList.add(new UserMessage(message));
        // 3. 控制历史消息数量，防止过大
        if (messageList.size() > MAX_HISTORY_SIZE) {
            // 保留系统消息，移除最旧的用户/AI消息
            messageList = messageList.subList(messageList.size() - MAX_HISTORY_SIZE, messageList.size());
        }
        // 4. 更新会话中的历史消息
        session.setAttribute(ChatController.CHAT_HISTORY_KEY, messageList);
        // 5. 创建包含历史的提示
        Prompt prompt = new Prompt(new ArrayList<>(messageList)); // 使用副本避免并发问题
        // 6. 生成AI响应并添加到历史
        List<Message> finalMessageList = messageList;
        Flux<String> sseFlux = chatModel.stream(prompt)
                .map(chatResponse -> {
                    // 提取AI回复内容
                    String content = chatResponse.getResult().getOutput().getText();
                    if (content != null && !content.isEmpty()) {
                        // 将AI回复添加到历史（注意：这里需要同步添加，Flux是异步的）
                        synchronized (finalMessageList) {
                            finalMessageList.add(new UserMessage(content));
                        }
                    }
//                    String sseContent = "data: " + (content != null ? content : "") + "\n\n";
                    return content;
//                    return "data: " + (content != null ? content : "") + "\n\n";
                });
        return sseFlux;
    }


    @Override
    public Flux<String> elkStreamChat(String question) {
        List<Document> relevantDocs = getRelevantDocuments(question);

        // 判断检索结果是否为空
        if (relevantDocs == null || relevantDocs.isEmpty()) {
            try {
                // 使用chatModel替代ollamaChatClient
                return chatModel.stream(new Prompt(new UserMessage(question)))
                        .map(ChatResponse::getResult)
                        .map(Generation::getOutput)  // 获取AssistantMessage
                        .map(AssistantMessage::getText)  // 从消息中获取内容
                        .flatMapSequential(Flux::just);
            } catch (Exception e) {
                throw new RuntimeException(e.getMessage());
            }
        } else {
            // 有检索结果 → 将文档作为上下文生成回答
            PromptTemplate promptTemplate = getPromptTemplate(relevantDocs, question);
            try {
                // 使用chatModel替代ollamaChatClient
                Prompt prompt = promptTemplate.create();
                return chatModel.stream(prompt)
                        .map(ChatResponse::getResult)
                        .map(Generation::getOutput)  // 获取AssistantMessage
                        .map(AssistantMessage::getText); // 从消息中获取内容
            } catch (Exception e) {
                throw new RuntimeException(e.getMessage());
            }
        }
    }

    @Override
    public String elkChatByMyClient(String question, HttpSession session){
        List<Document> relevantDocs = getRelevantDocuments(question);
        // 用sessionId作为当前会话的conversationId
        String conversationId = session.getId();
        //  判断检索结果是否为空
        if (relevantDocs==null || relevantDocs.isEmpty()) {
            try{
                return ollamaChatClient.prompt()
                        .advisors(a -> a.param("conversationId", conversationId)) // 动态传入
                        .user(question)
                        .call()
                        .content();
            }catch (Exception e){
                throw new RuntimeException(e.getMessage());
            }
        } else{
            // 有检索结果 → 将文档作为上下文生成回答
            PromptTemplate promptTemplate = getPromptTemplate(relevantDocs, question);
            try{
                return ollamaChatClient.prompt(promptTemplate.create())
                        .advisors(a -> a.param("conversationId", conversationId)) // 动态传入
                        .user(question)
                        .call()
                        .content();
            }catch (Exception e){
                throw new RuntimeException(e.getMessage());
            }

        }
    }

    @Override
    public Flux<String> elkStreamChatByMyClient(String question, HttpSession session) {
        List<Document> relevantDocs = getRelevantDocuments(question);
        // 新增日志：打印检索到的文档内容
        log.info("检索到的相关文档：{}", relevantDocs.stream()
                .map(Document::getText)
                .collect(Collectors.joining("\n")));
        // 用sessionId作为当前会话的conversationId
        String conversationId = session.getId();
        //  判断检索结果是否为空
        if (relevantDocs==null || relevantDocs.isEmpty()) {
            try{
                return  ollamaChatClient.prompt()
                        .user(question)
                        .advisors(a -> a.param("conversationId", conversationId)) // 动态传入
                        .stream()
                        .content()
                        .flatMapSequential(Flux::just);
            }catch (Exception e){
                throw new RuntimeException(e.getMessage());
            }
        }else {
            // 有检索结果 → 将文档作为上下文生成回答
            PromptTemplate promptTemplate = getPromptTemplate(relevantDocs, question);
            try{
                return ollamaChatClient.prompt(promptTemplate.create())
                        .advisors(a -> a.param("conversationId", conversationId)) // 动态传入
                        .user(question)
                        .stream()
                        .content();
            }catch (Exception e){
                throw new RuntimeException(e.getMessage());
            }
        }
    }

    private List<Document> getRelevantDocuments(String query) {
        // 从向量库检索相关文档（假设 topK=3）
        // 相似度阈值
        return vectorStore.similaritySearch(
                SearchRequest.builder()
                        .topK(3)
                        // 提高阈值，过滤低相关文档
                        .similarityThreshold(0.7)
                        .query(query)
                        .build()
        );
    }

    /*将文档作为上下文生成回答*/
    private PromptTemplate getPromptTemplate(List<Document> relevantDocs,String question) {
        String template = """
                1. 请严格按照以下知识库内容以及上下文回答问题，如果知识库与你的认知冲突，以知识库为准:
                {documents}
                
                2. 直接回答问题，不解释思考过程，不添加无关内容。
                3. 语言简洁，控制在最短篇幅内说清答案
                
                请基于上述知识库以及上下文，回答用户问题：：{question}
                """;
        PromptTemplate promptTemplate = new PromptTemplate(template);
        promptTemplate.add("documents", relevantDocs.stream()
                .map(Document::getText)
                .collect(Collectors.joining("\n")));
        promptTemplate.add("question", question);
        return promptTemplate;
    }

//    请严格按照以下步骤回答：
//            1. 只参考提供的【知识库内容】，完全忽略你之前的任何知识。
//            2. 如果知识库内容与你的认知冲突，以知识库为准。
//            3. 直接根据知识库回答问题，不要进行额外推理。
//
//            【知识库内容】
//    {documents}
//
//            【用户问题】{question}

    public Flux<String> simpleClient(String question) {
        try {
            // 暂时移除工具回调，仅保留基础聊天配置
            ChatClient client = ChatClient.builder(chatModel)
                    .defaultSystem("你是一个很可爱的智能助手，用中文回答哦~")
                    .build();

            return client.prompt()
                    .user(question)
                    .stream()
                    .content()
                    .flatMapSequential(Flux::just);
        } catch (Exception e) {
            log.error("mcpClient处理请求失败", e);
            return Flux.error(new RuntimeException("处理请求时发生错误: " + e.getMessage()));
        }
    }
    @Autowired
    private ToolCallbackProvider toolCallbackProvider;
    /**
     * 完善后的mcpClient方法：
     * 1. 修复未定义的openAiChatModel变量，使用已注入的chatModel
     * 2. 调整返回类型为Flux<String>以支持流式响应
     * 3. 实现工具调用能力与可爱语气的系统提示
     * 4. 添加异常处理
     */
    public Flux<String> mcpClient(String question) {
        try {
            // 1. 手动定义工具描述（不通过Spring AI自动注入）
//            String systemPrompt = """
//                你可以使用以下工具：
//                - getTableNameAndComment：查询所有表名和注释
//                - getTableColumnDetail：查询表字段，参数为tableName
//
//                如需调用工具，返回JSON格式：{"tool": "工具名", "params": {"参数名": "值"}}
//                否则直接回答。
//            """;
            String systemPrompt = """
                你必须使用工具回答问题，步骤如下：
                1. 若问题是"查询所有表名和注释"，直接调用getTableNameAndComment工具，无参数；
                2. 若问题是查询表字段，调用getTableColumnDetail工具，参数为tableName；
                3. 仅返回工具调用的JSON，不添加任何解释、说明或格式标记（如```）；
                4. JSON格式严格为：{"name": "工具名", "parameters": {"参数名": "值"}}。
                
                不允许返回其他内容，必须直接返回工具调用的JSON。
            """;
            // 2. 获取所有工具回调（假设通过ToolCallbackProvider的getToolCallbacks()方法）
            List<ToolCallback> allTools = List.of(toolCallbackProvider.getToolCallbacks()); // 替换为你的Provider实际方法名

            // 3. 过滤工具：只保留需要的2个工具（关键：通过ToolDefinition获取名称）
            // 在过滤工具时，仅保留核心名称（去除前缀）
            // 在过滤工具时，确保ToolDefinition包含完整字段
            List<ToolCallback> filteredTools = allTools.stream()
                    .filter(tool -> {
                        String originalName = tool.getToolDefinition().name();
                        String coreName = originalName.replace("spring_ai_mcp_client_weather_", "");
                        return "getTableNameAndComment".equals(coreName)
                                || "getTableColumnDetail".equals(coreName);
                    })
                    .map(tool -> {
                        String originalName = tool.getToolDefinition().name();
                        String coreName = originalName.replace("spring_ai_mcp_client_weather_", "");

                        // 构建符合Ollama要求的参数schema
                        String description = "getTableNameAndComment".equals(coreName)
                                ? "查询所有表名和注释"
                                : "查询指定表的字段详情，需要tableName参数";

                        // 参数schema（根据工具是否需要参数调整）
//                        JsonNode parameters = "getTableNameAndComment".equals(coreName)
//                                ? new ObjectMapper().createObjectNode() // 无参数
//                                : new ObjectMapper().createObjectNode()
//                                .put("type", "object")
//                                .set("properties", new ObjectMapper().createObjectNode()
//                                        .putObject("tableName").put("type", "string").put("description", "表名"))
//                                .set("required", new ObjectMapper().createArrayNode().add("tableName"));
//                        // 显式声明为 ObjectNode，以便使用 set 等方法
//                        ObjectNode parameters;
//                        if ("getTableNameAndComment".equals(coreName)) {
//                            parameters = new ObjectMapper().createObjectNode(); // 无参数，创建空的 ObjectNode
//                        } else {
//                            parameters = new ObjectMapper().createObjectNode();
//                            parameters.put("type", "object");
//                            // 创建 properties 节点并设置 tableName 的属性
//                            ObjectNode propertiesNode = new ObjectMapper().createObjectNode();
//                            ObjectNode tableNameNode = propertiesNode.putObject("tableName");
//                            tableNameNode.put("type", "string");
//                            tableNameNode.put("description", "表名");
//                            parameters.set("properties", propertiesNode);
//                            // 创建 required 数组并添加 tableName
//                            ArrayNode requiredNode = new ObjectMapper().createArrayNode();
//                            requiredNode.add("tableName");
//                            parameters.set("required", requiredNode);
//                        }


                        ToolDefinition simplifiedDefinition = ToolDefinition.builder()
                                .name(coreName)
                                .description(description)
                                //.inputSchema(parameters.toString()) // 必须包含参数schema
                                .inputSchema(buildToolSchema(coreName)) // 抽离为独立方法，确保格式正确
                                .build();

                        return new ToolCallback() {
                            @Override
                            public ToolDefinition getToolDefinition() { return simplifiedDefinition; }
                            @Override
                            public String call(String toolInput) { return tool.call(toolInput); }
                        };
                    })
                    .collect(Collectors.toList());

             //4. 验证过滤结果（调试用）
            log.info("过滤后保留的工具: {}", filteredTools.stream()
                    .map(t -> t.getToolDefinition().name())
                    .collect(Collectors.toList()));
            // 3. 构建客户端时使用过滤后的工具回调
            // 构建具有工具调用能力的聊天客户端
            // 关键修复：将defaultTools()改为defaultToolCallbacks()
            //Ollama 在 0.1.46 版本后才支持tools参数，低于此版本会直接返回 400；
//            ChatClient client = ChatClient.builder(chatModel)
//                    .defaultToolCallbacks(callbackProvider)  // ToolCallbackProvider callbackProvider
//                    .defaultSystem(systemPrompt)
//                    .build();
            ChatClient client = ChatClient.builder(chatModel)
                    .defaultToolCallbacks(filteredTools)  // 仅传入需要的工具
                    .defaultSystem(systemPrompt)
                    .build();

            // 流式获取响应并返回
            return client.prompt()
                    .user(question)
                    .stream()
                    .content()
                    .flatMapSequential(Flux::just);
        } catch (Exception e) {
            log.error("mcpClient处理请求失败", e);
            return Flux.error(new RuntimeException("处理请求时发生错误: " + e.getMessage()));
        }
    }



    // 独立方法：构建符合Ollama要求的工具Schema
    private String buildToolSchema(String coreName) {
//                        {
//                            "tools": [
//                                {
//                                    "name": "getTableNameAndComment",
//                                    "description": "查询所有表名和注释",
//                                    "parameters": {
//                                        "type": "object",
//                                        "properties": {}, // 无参数时为空对象
//                                        "required": []
//                                    }
//                                }
//                            }
//                          ]
//                        }

        ObjectMapper mapper = new ObjectMapper();
        ObjectNode schema = mapper.createObjectNode();
        schema.put("type", "object"); // 顶层必须声明为object

        if ("getTableNameAndComment".equals(coreName)) {
            // 无参数工具：properties为空对象，required为空数组
            schema.set("properties", mapper.createObjectNode());
            schema.set("required", mapper.createArrayNode());
        } else {
            // 有参数工具：定义tableName的属性
            ObjectNode properties = mapper.createObjectNode();
            ObjectNode tableNameProp = properties.putObject("tableName");
            tableNameProp.put("type", "string");
            tableNameProp.put("description", "表名");

            ArrayNode required = mapper.createArrayNode();
            required.add("tableName");

            schema.set("properties", properties);
            schema.set("required", required);
        }
        return schema.toString();
    }

}
